package com.bwie.ai.controller;

import com.bwie.ai.advisor.LoggingAdvisor;
import org.springframework.ai.chat.client.ChatClient;
import org.springframework.ai.chat.client.advisor.AbstractChatMemoryAdvisor;
import org.springframework.ai.chat.client.advisor.PromptChatMemoryAdvisor;
import org.springframework.ai.chat.client.advisor.QuestionAnswerAdvisor;
import org.springframework.ai.chat.memory.ChatMemory;
import org.springframework.ai.vectorstore.SearchRequest;
import org.springframework.ai.vectorstore.VectorStore;
import org.springframework.beans.factory.annotation.Autowired;
import org.springframework.web.bind.annotation.GetMapping;
import org.springframework.web.bind.annotation.RequestMapping;
import org.springframework.web.bind.annotation.RequestParam;
import org.springframework.web.bind.annotation.RestController;
import reactor.core.publisher.Flux;
import org.springframework.http.MediaType;
import java.time.LocalDate;

@RestController
@RequestMapping("/ai")
public class OpenAiController {
    private final ChatClient chatClient;

    public OpenAiController(ChatClient.Builder chatClientBuilder, ChatMemory chatMemory, VectorStore vectorStore) {
        this.chatClient = chatClientBuilder
                .defaultSystem("""
					   你是"小李今天怎么样"医疗系统的客户聊天支持代理，请以友好的方式来回复，
                                                                                       你可以通过在线聊天系统与客户互动，在提供了有关病症的信息之前，不过不能回答除了医疗方面以外的问题，
                                                                                                          您必须始终从用户获取以下信息：病状。
                                                                                                          在询问患者之前，你需要检查消息历史记录以获取此信息。
                                                                                                          你可以根据患者的病状高效地为患者匹配到相对应领域的医生。
                                                                                                          在进行挂号和取消挂号之前,需要获取医生姓名和患者姓名。
                                                                                                          取消挂号前根据医生姓名和患者姓名获取挂号状态，不需要显示状态代码。
                                                                                                          今天的日期是 {current_date}.
					""")
                .defaultAdvisors(
                        new PromptChatMemoryAdvisor(chatMemory),
                        new LoggingAdvisor(),
                        new QuestionAnswerAdvisor(vectorStore, SearchRequest.defaults())
                )
                .defaultFunctions("getDoctorBySymptom","register","cancelRegister","getRegisterState")
                .build();
    }

    @GetMapping(value = "/chat",produces = MediaType.TEXT_EVENT_STREAM_VALUE)
    public Flux<String> generation(@RequestParam(value = "message") String message) {
        Flux<String> content = this.chatClient.prompt()
                .user(message)
                .system(promptSystemSpec -> promptSystemSpec.param("current_date", LocalDate.now().toString()))
                .advisors(advisorSpec -> advisorSpec.param(AbstractChatMemoryAdvisor.CHAT_MEMORY_RETRIEVE_SIZE_KEY,100))
                .stream().content();
        return content.concatWith(Flux.just("[complete]"));
    }
}
